﻿//
// Created by QU on 24-10-21.
//

#include <exception>
#include <iostream>
#include <format>
#include <chrono>
#include <thread>
#include <span>

#include <boost/asio/co_spawn.hpp>
#include <boost/asio/detached.hpp>
#include <boost/asio.hpp>

namespace asio = boost::asio;
namespace this_coro = asio::this_coro;

using asio::awaitable;
using asio::co_spawn;
using asio::detached;
using asio::use_awaitable;

using tcp = asio::ip::tcp;

// 自定义格式化器为 boost::asio::ip::tcp::endpoint
template <>
struct std::formatter<boost::asio::ip::tcp::endpoint> : std::formatter<std::string> {
    auto format(const boost::asio::ip::tcp::endpoint& ep, format_context& ctx) const {
        return formatter<string>::format(ep.address().to_string() + ":" + std::to_string(ep.port()), ctx);
    }
};

// 辅助函数：获取当前时间戳
std::string get_timestamp() {
    auto now = std::chrono::system_clock::now();
    return std::format("{:%Y-%m-%d %H:%M:%S}", now);
}

// 辅助函数：获取当前线程ID
std::string get_thread_id() {
    return std::format("{}", std::this_thread::get_id());
}

// 日志输出函数
template<typename... Args>
void log(std::string_view fmt, const Args&... args) {
    auto formatted = std::vformat(fmt, std::make_format_args(args...));
    std::cout << std::format("[{}][Thread {}] {}\n",
                             get_timestamp(), get_thread_id(), formatted);
}

awaitable<void> echo(tcp::socket socket)
{
    try
    {
        auto remote_ep = socket.remote_endpoint();
        log("Echo coroutine started for client: {}", remote_ep);

        std::array<std::byte, 1024> buffer;
        for(;;)
        {
            log("Waiting for data from client: {}", remote_ep);

            std::size_t size = co_await socket.async_read_some(asio::buffer(buffer), use_awaitable);

            log("Received {} bytes from client: {}", size, remote_ep);
            log("Echoing data back to client: {}", remote_ep);

            co_await async_write(socket, asio::buffer(std::span(buffer).first(size)), use_awaitable);

            log("Data echoed back to client: {}", remote_ep);
        }
    }
    catch (std::exception& e)
    {
        log("Exception in echo coroutine: {}", e.what());
    }
}

awaitable<void> listener()
{
    auto executor = co_await this_coro::executor;
    tcp::acceptor acceptor(executor, {tcp::v4(), 10120});

    log("Listener started on port 10120");

    for (;;)
    {
        log("Waiting for new connection");

        tcp::socket socket = co_await acceptor.async_accept(use_awaitable);
        auto remote_ep = socket.remote_endpoint();

        log("New connection accepted from: {}", remote_ep);

        co_spawn(executor, echo(std::move(socket)), detached);

        log("Spawned new echo coroutine for client: {}", remote_ep);
    }
}

int main()
{
    try
    {
        asio::io_context io_context(1);
        asio::signal_set signals(io_context, SIGINT, SIGTERM);
        signals.async_wait([&](auto, auto)
        {
            log("Received stop signal, stopping io_context");
            io_context.stop();
        });

        log("Starting listener coroutine");
        co_spawn(io_context, listener, detached);

        log("Running io_context");
        io_context.run();

        log("io_context stopped");
    }
    catch (std::exception& e)
    {
        log("Exception in main: {}", e.what());
    }
}
